import os
import time
from typing import List, Union, Dict
from openai import OpenAI
from openai import APIError, APITimeoutError, RateLimitError


def text_embeddings(structured_data: Union[Dict, List[Dict]]) -> List[List[float]]:
    """
    从结构化元数据生成中文法律文本的向量表示，适配劳动合同法等法律条文

    参数:
        structured_data: 单个或列表形式的结构化元数据字典，需包含'vector_text'字段
                         即上一步骤structure_labor_law_metadata函数的输出结果

    返回:
        向量列表，每个元素为对应文本的1024维向量，与输入数据顺序一致
    """
    # 初始化客户端

    print("structured_data:::::::",structured_data)
    api_key = "sk-efa04219920b4938b7573cf1b80f9784"
    if not api_key:
        raise ValueError("请配置DASHSCOPE_API_KEY环境变量或直接替换api_key参数")

    client = OpenAI(
        api_key=api_key,
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

    # 统一转换为列表格式
    if isinstance(structured_data, dict):
        structured_data = [structured_data]

    # 验证输入数据格式
    for i, item in enumerate(structured_data):
        if "vector_text" not in item:
            raise ValueError(f"结构化数据第{i}项缺少'vector_text'字段")

    # 提取需要生成向量的文本
    texts = [item["vector_text"] for item in structured_data]

    # 法律文本预处理（增强法律术语权重）
    processed_texts = [_preprocess_legal_text(text) for text in texts]

    print("processed_texts:::::::",processed_texts)

    # 百炼text-embedding-v4的最佳实践参数
    batch_size = 10  # 模型最大批处理限制
    embeddings = []
    retries = 3  # 重试次数
    retry_delay = 2  # 重试延迟（秒）

    # 分批处理
    for i in range(0, len(processed_texts), batch_size):
        batch = processed_texts[i:i + batch_size]
        print(f"处理批次：{batch}")
        batch_indices = list(range(i, min(i + batch_size, len(processed_texts))))

        for attempt in range(retries):
            try:
                print("开始处理批次：", batch)
                response = client.embeddings.create(
                    model="text-embedding-v4",
                    input=batch,
                    dimensions=1024,  # 法律文本推荐1024维（平衡精度与效率）
                    encoding_format="float",
                    # 法律领域提示词（增强专业相关性）
                    user="legal_rag_system"
                )

                print("responseresponseresponseresponseresponseresponse*****:",response)

                # 验证返回数量是否匹配
                if len(response.data) != len(batch):
                    raise ValueError(f"批量处理失败：输入{len(batch)}条，返回{len(response.data)}条")

                # 提取向量并保持原始顺序
                batch_embeddings = [item.embedding for item in response.data]
                embeddings.extend(batch_embeddings)
                break  # 成功则跳出重试循环

            except (APIError, APITimeoutError) as e:
                print(f"处理批次{batch_indices}时出错：{str(e)}")
                if attempt < retries - 1:
                    time.sleep(retry_delay * (attempt + 1))  # 指数退避
                    continue
                raise RuntimeError(f"处理批次{batch_indices}失败：{str(e)}")
            except RateLimitError:
                if attempt < retries - 1:
                    time.sleep(retry_delay * (attempt + 1) * 2)  # 限流时延长等待
                    continue
                raise RuntimeError(f"API限流，处理批次{batch_indices}失败")

    return embeddings


def _preprocess_legal_text(text: str) -> str:
    """
    法律文本预处理，增强向量对法律术语的敏感性
    """
    # 移除多余空格和换行
    processed = ' '.join(text.split())

    # 增强法律主体术语（用人单位、劳动者、劳动合同等）
    legal_terms = {
        "用人单位": "用人单位[法律主体]",
        "劳动者": "劳动者[法律主体]",
        "劳动合同": "劳动合同[法律客体]",
        "经济补偿": "经济补偿[法律责任]",
        "赔偿金": "赔偿金[法律责任]",
        "试用期": "试用期[合同条款]",
        "解除": "解除[合同行为]",
        "终止": "终止[合同行为]"
    }

    # 替换术语以增强权重
    for term, enhanced in legal_terms.items():
        processed = processed.replace(term, enhanced)

    return processed


# 测试函数
if __name__ == "__main__":


    # query_emb = text_embeddings([{"vector_text": "未签劳动合同的补偿"}])[0]
    # print(f"向量维度：{query_emb}")

    #模拟上一步骤生成的结构化数据
    sample_structured_data = [
        {
            "id": "labor_law_010",
            "vector_text": "第十条 建立劳动关系，应当订立书面劳动合同。已建立劳动关系，未同时订立书面劳动合同的，应当自用工之日起一个月内订立书面劳动合同。用人单位与劳动者在用工前订立劳动合同的，劳动关系自用工之日起建立。",
            "metadata": {
                "law_name": "中华人民共和国劳动合同法",
                "chapter": "第二章 劳动合同的订立",
                "article_num": "第十条",
                # 其他元数据字段...
            }
        },
        {
            "id": "labor_law_046",
            "vector_text": "第四十六条 有下列情形之一的，用人单位应当向劳动者支付经济补偿：（一）劳动者依照本法第三十八条规定解除劳动合同的；（二）用人单位依照本法第三十六条规定向劳动者提出解除劳动合同并与劳动者协商一致解除劳动合同的；",
            "metadata": {
                "law_name": "中华人民共和国劳动合同法",
                "chapter": "第四章 劳动合同的解除和终止",
                "article_num": "第四十六条",
                # 其他元数据字段...
            }
        }
    ]

    try:
        embeddings = text_embeddings(sample_structured_data)
        print(f"生成向量数量：{len(embeddings)}")
        print(f"第一条向量维度：{len(embeddings[0])}")
        print(f"向量示例（前5位）：{embeddings[0][:5]}")
    except Exception as e:
        print(f"测试失败：{str(e)}")
